"""
reference https://www.youtube.com/watch?v=h9gpufJFF-0
"""

import numpy as np
import pandas as pd
import itertools

np.random.seed(1)

uu_data = {
    "u1": {"hp1": 4, "tw": 5, "sw1": 1},
    "u2": {"hp1": 5, "hp2": 5, "hp3": 4},
    "u3": {"tw": 2, "sw1": 4, "sw2": 5},
    "u4": {"hp2": 3, "sw3": 3},
}


ii_data = {
    "m1": {"1": 1, "3": 3, "6": 5, "9": 5, "11": 4},
    "m2": {"3": 5, "4": 4, "7": 4, "10": 2, "11": 1, "12": 3},
    "m3": {"1": 2, "2": 4, "4": 1, "5": 2, "7": 3, "9": 4, "10": 3, "11": 5},
    "m4": {"2": 2, "3": 4, "5": 5, "8": 4, "11": 2},
    "m5": {"3": 4, "4": 3, "5": 4, "6": 2, "11": 2, "12": 5},
    "m6": {"1": 1, "3": 3, "5": 3, "8": 2, "11": 4}
}


def pad_zero2unrated(rating_dict, sort_col=False):
    res_dict = {}
    col = list(set(itertools.chain(*[list(v.keys()) for v in rating_dict.values()])))
    if sort_col: col.sort()
    for k, rating in rating_dict.items():
        res_dict[k] = [rating_dict[k].get(c, 0) for c in col]
    res_tb = pd.DataFrame.from_dict(res_dict, orient="index", dtype=np.float32)
    res_tb.columns = col
    return res_tb


def centred_cosine(selected, others):
    """
    (R1*R2)/(||R1||*||R2||)
    subtract the rating by the mean of this user
    :param user_rating: rates for all movies
    :param others_rating: rates for all movies
    :return:
    """
    assert others.ndim == 2
    if selected.ndim == 1:
        scaled_selected_rating = selected[None, :].copy()
    else:
        scaled_selected_rating = selected.values.copy()
    scaled_others_rating = others.values.copy()
    rows = np.concatenate((scaled_selected_rating, scaled_others_rating), axis=0)
    for row in rows:
        non_zero_idx = np.argwhere(row != 0).ravel()
        non_zero = row[non_zero_idx]
        v = non_zero - non_zero.mean()
        row[non_zero_idx] = v
    scaled_selected_rating = rows[0, :]
    scaled_others_rating = rows[1:, :]

    # cosine similarity
    r1, r2 = scaled_selected_rating, scaled_others_rating
    numerator = r1.dot(r2.T).ravel()
    denominator = np.linalg.norm(r1) * np.linalg.norm(r2, axis=1)
    return np.divide(numerator, denominator, where=denominator != 0)


def find_n_similar(selected, others, n=None):
    similarity = centred_cosine(selected, others)
    if n is None:
        top_n_idx = np.arange(len(others))
    else:
        top_n_idx = np.argsort(similarity)[-n:]
    top_n_name = others.index[top_n_idx]
    return pd.Series(similarity[top_n_idx], index=top_n_name)


def user2userCF(user_name, rating_dict, user_as_key):
    rating_tb = pad_zero2unrated(rating_dict)
    if not user_as_key:    # make user at the first dim
        rating_tb = rating_tb.T

    user_rating = rating_tb.loc[user_name]
    other_users_name = list(rating_tb.index)
    other_users_name.remove(user_name)
    others_rating = rating_tb.loc[other_users_name]

    top_n_user_similarity = find_n_similar(user_rating, others_rating, 2)
    unwatched_movies = user_rating[user_rating == 0].index

    base_rating = others_rating.loc[top_n_user_similarity.index, unwatched_movies]
    pred_rating = (base_rating.values * top_n_user_similarity.values[:, None]).sum(axis=0)/top_n_user_similarity.values.sum()
    return pd.Series(pred_rating, base_rating.columns)


def item2itemCF(user_name, rating_dict, user_as_key):
    """
    item2item is generally better than user2user in practice
    :param user_name:
    :param rating_data:
    :return:
    """
    rating_tb = pad_zero2unrated(rating_dict)
    if user_as_key:        # make item at the first dim
        rating_tb = rating_tb.T

    unwatched_movies = rating_tb[user_name][rating_tb[user_name] == 0].index
    pred_rating = {m: 0 for m in unwatched_movies}
    for movie_name in unwatched_movies:
        # for each movie, find similar movies
        movie_rating = rating_tb.loc[movie_name]
        other_movies_name = list(rating_tb.index)
        other_movies_name.remove(movie_name)
        others_rating = rating_tb.loc[other_movies_name]
        top_n_movie_similarity = find_n_similar(movie_rating, others_rating, 2)

        # find user, calculate rate based on the similarity and rating history of those similar movies
        base_rating = others_rating.loc[top_n_movie_similarity.index, user_name]
        pred_rating[movie_name] = (top_n_movie_similarity*base_rating).sum()/top_n_movie_similarity.sum()
    return pd.Series(pred_rating,).astype(np.float32)


print(user2userCF("5", ii_data, user_as_key=False))
print(item2itemCF("5", ii_data, user_as_key=False))